#pragma once
#include <string>
#include <ncnn/net.h>
#include <opencv2/opencv.hpp>

namespace onechchy
{
	enum class PromptType
	{
		Point = 0,
		Box
	};

	typedef struct _ImageInfo
	{
		int imgW;
		int imgH;
		int padW;
		int padH;
		float scale;
	}ImageInfo;

	typedef struct _PromptInfo
	{
		std::vector<float> points;
		std::vector<int> labels;
		cv::Mat mask;
		int hasMask;
		PromptType promptType;
	}PromptInfo;

	typedef struct _SamResult
	{
		cv::Mat mask;
		cv::Rect box;
		float iouPred;
		float stableScore;
	}SamResultT;

	typedef struct _PipelineResult
	{
		ncnn::Mat imageEmbeddings;
		ImageInfo imageInfo;
		PromptInfo promptInfo;
		std::vector<SamResultT> samResults;
	}PipeLineResult;

	class SegmentAnything
	{
	public:
		SegmentAnything() = default;
		~SegmentAnything();

		int load(const std::string& imageEncoderParam, const std::string& imageEncoderBin, const std::string& maskDecoderParam, const std::string& maskDecoderBin);
		int imageEncoder(const cv::Mat& bgr, ncnn::Mat& imageEmbeddings, ImageInfo& imageInfo);
		int maskDecoder(const ncnn::Mat& imageEmbeddings, ImageInfo& imageInfo, const PromptInfo& promptInfo, std::vector<SamResultT>& samResult, float predIouThresh = 0.88f, float stabilityScoreThresh = 0.95f);
		int nms(const cv::Mat& bgr, std::vector<SamResultT>& proposals, std::vector<int>& picked, float nmsThreshold = 0.7f);

	private:
		int embedPoints(const PromptInfo& promptInfo, std::vector<ncnn::Mat>& pointLabels, ncnn::Mat& pointCoords);
		int transformCoords(const ImageInfo& imageInfo, ncnn::Mat& pointCoords);
		int embedMasks(const PromptInfo& promptInfo, ncnn::Mat& maskInput, ncnn::Mat& hasMask);
		float calculateStabilityScore(cv::Mat& mask, float maskThreshold = 0.f, float stableScoreOffset = 1.f);
		int postprocessMask(cv::Mat& mask, cv::Rect& box);

	private:
		const float means[3] = { 123.675f, 116.28f, 103.53f };
		const float norms[3] = { 0.01712475f, 0.0175f, 0.01742919f };
		ncnn::Net mImageEncoderNet;
		ncnn::Net mMaskDecoderNet;
	};

	class PipeLine
	{
	public:
		PipeLine() = default;
		~PipeLine();

		int init(const std::string& imageEncoderParam, const std::string& imageEncoderBin, const std::string& maskDecoderParam, const std::string& maskDecoderBin);
		int imageEmbedding(const cv::Mat& bgr, PipeLineResult& pipelineResult);
		int predict(const cv::Mat& bgr, PipeLineResult& pipelineResult);
		int autoPredict(const cv::Mat& bgr, PipeLineResult& pipelineResult, int nPerSide = 32);
		void draw(const cv::Mat& bgr, const PipeLineResult& pipelineResult);

	private:
		void getGridPoints(std::vector<float>& pointsXyVec, int nPerSide);
		std::shared_ptr<SegmentAnything> mSam;
	};
}

